{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install autogen-agentchat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install autogen-ext[azure,openai]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "from enum import Enum\n",
    "from typing import List, Optional, Union\n",
    "\n",
    "class AgentEnum(str, Enum):\n",
    "    FlightBooking = \"flight_booking\"\n",
    "    HotelBooking = \"hotel_booking\"\n",
    "    CarRental = \"car_rental\"\n",
    "    ActivitiesBooking = \"activities_booking\"\n",
    "    DestinationInfo = \"destination_info\"\n",
    "    DefaultAgent = \"default_agent\"\n",
    "    GroupChatManager = \"group_chat_manager\"\n",
    "\n",
    "# Travel SubTask Model\n",
    "class TravelSubTask(BaseModel):\n",
    "    task_details: str\n",
    "    assigned_agent: AgentEnum # we want to assign the task to the agent\n",
    "\n",
    "class TravelPlan(BaseModel):\n",
    "    main_task: str\n",
    "    subtasks: List[TravelSubTask]\n",
    "    is_greeting: bool\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set the environment variable\n",
    "os.environ[\"AZURE_OPENAI_ENDPOINT\"] = \"https://AZURE_OPENAI_ENDPOINT.openai.azure.com/\"\n",
    "os.environ[\"AZURE_OPENAI_API_KEY\"] = \"AZURE_OPENAI_API_KEY\"\n",
    "os.environ[\"AZURE_OPENAI_DEPLOYMENT_NAME\"] = \"gpt-4o-2024-08-06\"\n",
    "os.environ[\"AZURE_OPENAI_API_VERSION\"] = \"2024-08-01-preview\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from typing import Optional\n",
    "\n",
    "from autogen_core.models import UserMessage, SystemMessage, AssistantMessage\n",
    "from autogen_ext.models.openai import AzureOpenAIChatCompletionClient\n",
    "\n",
    "\n",
    "# Function to get environment variable and ensure it is not None\n",
    "def get_env_variable(name: str) -> str:\n",
    "    value = os.getenv(name)\n",
    "    if value is None:\n",
    "        raise ValueError(f\"Environment variable {name} is not set\")\n",
    "    return value\n",
    "\n",
    "\n",
    "# Create the client with type-checked environment variables\n",
    "client = AzureOpenAIChatCompletionClient(\n",
    "    azure_deployment=get_env_variable(\"AZURE_OPENAI_DEPLOYMENT_NAME\"),\n",
    "    model=get_env_variable(\"AZURE_OPENAI_DEPLOYMENT_NAME\"),\n",
    "    api_version=get_env_variable(\"AZURE_OPENAI_API_VERSION\"),\n",
    "    azure_endpoint=get_env_variable(\"AZURE_OPENAI_ENDPOINT\"),\n",
    "    api_key=get_env_variable(\"AZURE_OPENAI_API_KEY\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'is_greeting': False,\n",
      " 'main_task': 'Plan a family trip from Singapore to Melbourne.',\n",
      " 'subtasks': [{'assigned_agent': 'flight_booking',\n",
      "               'task_details': 'Book round-trip flights from Singapore to '\n",
      "                               'Melbourne.'},\n",
      "              {'assigned_agent': 'hotel_booking',\n",
      "               'task_details': 'Find family-friendly hotels in Melbourne.'},\n",
      "              {'assigned_agent': 'car_rental',\n",
      "               'task_details': 'Arrange a car rental suitable for a family of '\n",
      "                               'four in Melbourne.'},\n",
      "              {'assigned_agent': 'activities_booking',\n",
      "               'task_details': 'List family-friendly activities in Melbourne.'},\n",
      "              {'assigned_agent': 'destination_info',\n",
      "               'task_details': 'Provide information about Melbourne as a '\n",
      "                               'travel destination.'}]}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_28718/3425863192.py:17: UserWarning: Resolved model mismatch: gpt-4o-2024-08-06 != gpt-4o-2024-11-20. Model mapping in autogen_ext.models.openai may be incorrect.\n",
      "  response = await client.create(messages=messages, extra_create_args={\"response_format\": TravelPlan})\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "# Define the user message\n",
    "messages = [\n",
    "    SystemMessage(content=\"\"\"You are an planner agent.\n",
    "    Your job is to decide which agents to run based on the user's request.\n",
    "    Below are the available agents specialised in different tasks:\n",
    "    - FlightBooking: For booking flights and providing flight information\n",
    "    - HotelBooking: For booking hotels and providing hotel information\n",
    "    - CarRental: For booking cars and providing car rental information\n",
    "    - ActivitiesBooking: For booking activities and providing activity information\n",
    "    - DestinationInfo: For providing information about destinations\n",
    "    - DefaultAgent: For handling general requests\"\"\", source=\"system\"),\n",
    "    UserMessage(content=\"Create a travel plan for a family of 2 kids from Singapore to Melboune\", source=\"user\"),\n",
    "]\n",
    "\n",
    "\n",
    "response = await client.create(messages=messages, extra_create_args={\"response_format\": TravelPlan})\n",
    "\n",
    "# Ensure the response content is a valid JSON string before loading it\n",
    "response_content: Optional[str] = response.content if isinstance(response.content, str) else None\n",
    "if response_content is None:\n",
    "    raise ValueError(\"Response content is not a valid JSON string\")\n",
    "\n",
    "# Print the response content after loading it as JSON\n",
    "pprint(json.loads(response_content))\n",
    "\n",
    "# Validate the response content with the MathReasoning model\n",
    "# TravelPlan.model_validate(json.loads(response_content))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
